package com.arch.flink.biz;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.builder.Tuple2Builder;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

/**
 * @author pizhihui
 * @date 2020-08-07
 */
public class TestKeyedStateMain {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        Tuple2Builder<Long, Long> builder = new Tuple2Builder<>();
        builder.add(1L, 3L);
        builder.add(1L, 5L);
        builder.add(1L, 7L);

        builder.add(2L, 4L);
        builder.add(2L, 2L);
        builder.add(2L, 5L);

        Tuple2<Long, Long>[] tuple2s = builder.build();
        DataStreamSource<Tuple2<Long, Long>> dataStreamSource = env.fromElements(tuple2s);

//        DataStreamSource<Tuple2<Long, Long>> dataStreamSource = env.fromElements(
//                Tuple2.of(1L, 3L), Tuple2.of(1L, 5L), Tuple2.of(1L, 7L),
//                Tuple2.of(2L, 4L), Tuple2.of(2L, 2L), Tuple2.of(2L, 5L));

        dataStreamSource.keyBy(0)
                .flatMap(new CountWindowAverageWithValueState())
                .print();

        env.execute("statful api test");

    }

}

/**
 * ValueState<T> ：这个状态为每一个 key 保存一个值
 * value() 获取状态值
 * update() 更新状态值
 * clear() 清除状态
 */
class CountWindowAverageWithValueState extends RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Double>> {

    // 用以保存每个 key 出现的次数，以及这个 key 对应的 value 的总值
    // managed keyed state
    // 1. ValueState 保存的是对应的一个 key 的一个状态值
    private ValueState<Tuple2<Long, Long>> countAndSum;

    @Override
    public void open(Configuration parameters) throws Exception {
        ValueStateDescriptor<Tuple2<Long, Long>> stateDescriptor =
                new ValueStateDescriptor<>("average", Types.TUPLE(Types.LONG, Types.LONG));

        countAndSum = getRuntimeContext().getState(stateDescriptor);
    }

    @Override
    public void flatMap(Tuple2<Long, Long> element, Collector<Tuple2<Long, Double>> collector) throws Exception {

        // 拿到当前的 key 的状态值
        Tuple2<Long, Long> currentState = countAndSum.value();
        // 如果状态值还没有初始化，则初始化
        if (null == currentState) {
            currentState = Tuple2.of(0L, 0L);
        }
        // 更新状态值中的元素的个数
        currentState.f0 += 1;

        // 更新状态值中的总值
        currentState.f1 += element.f1;

        // 更新状态
        countAndSum.update(currentState);


        // 判断，如果当前的 key 出现了 3 次，则需要计算平均值，并且输出
        if (currentState.f0 >= 3) {
            double avg = (double) currentState.f1 / currentState.f0;
            // 输出 key 及其对应的平均值
            collector.collect(Tuple2.of(element.f0, avg));
            //  清空状态值
            countAndSum.clear();
        }

    }
}
